package org.ala.tiktools.sharddb.mybatis.executor;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.mybatis.pager.Pagement;
import org.ala.tiktools.mybatis.pager.Pager;
import org.ala.tiktools.mybatis.pager.PagerHandler;
import org.ala.tiktools.sharddb.db.ThreadLocalLookupKey;
import org.ala.tiktools.sharddb.exception.ShardDBRuntimeException;
import org.ala.tiktools.sharddb.mybatis.interceptor.Shard;
import org.ala.tiktools.sharddb.mybatis.interceptor.ShardInfoThreadLocal;
import org.ala.tiktools.sharddb.mybatis.pager.ShardPagerHandler;
import org.ala.tiktools.sharddb.mybatis.router.ShardInfo;
import org.ala.tiktools.tools.CachingKVTools;
import org.slf4j.Logger;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * 非事务执行模式
 *
 * @author ala
 * @date 2025-02-24 14:01
 */
public class ShardNoTxExecutor {


    static Logger log = LoggerFactory.getLogger("shard", "SqlSession");


    /**
     *  分页处理
     */
    protected ShardPagerHandler pagerHandler = new ShardPagerHandler();


    /**
     *  执行非事务方法
     */
    protected Map<ShardInfo, ShardExecutor.Result> doShardInvokeWithoutTx(Shard anno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method, Object[] args) {
        //  取分页注解
        Pagement pagement = pagementMapper.get(method);

        //  不含分页的处理方式
        if (pagement == null) {
            return doInvoke(anno, threadPool, shardInfoMap, proxy, method);
        }
        //  含分页的处理方式
        else {
            return pagerHandler.doPage(anno, pagement, threadPool, proxy, method, args, shardInfoMap);
        }
    }


    /**
     *  非事务执行方式
     */
    public static Map<ShardInfo, ShardExecutor.Result> doInvoke(Shard anno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method) {
        return doInvoke(anno, threadPool, shardInfoMap, proxy, method, (a, s, r, o, m) -> {
            return ShardExecutor.doInvoke(a, s, m, r);
        });
    }
    /**
     *  非事务执行方式
     */
    public static Map<ShardInfo, ShardExecutor.Result> doInvoke(Shard anno, ThreadPoolExecutor threadPool, Map<ShardInfo, Object[]> shardInfoMap, Object proxy, Method method, IShardNoTxDoInvoker invoker) {
        Map<ShardInfo, ShardExecutor.Result> res = new ConcurrentHashMap<>();
        CountDownLatch latch = new CountDownLatch(shardInfoMap.size());

        long now = System.currentTimeMillis();
        shardInfoMap.forEach((shardInfo, shardArgs) -> {
            long _now = System.currentTimeMillis();
            try {
                threadPool.execute(() -> {
                    ThreadLocalLookupKey.setLookup(shardInfo.getShardDB());
//                    ShardInfoThreadLocal.set(shardInfo);
                    try {
                        ShardExecutor.Result r = invoker.doInvoke(anno, shardInfo, shardArgs, proxy, method);
//                        ShardExecutor. Result r = ShardExecutor.doInvoke(anno, shardInfo.getShardDB(), shardInfo.getShardTable(), proxy, method, shardArgs);
                        res.put(shardInfo, r.timeOfExecute(System.currentTimeMillis() - now));
                    } catch (Throwable e) {
                        res.put(shardInfo, ShardExecutor.Result.build(shardInfo, false, shardArgs, null, e, System.currentTimeMillis() - now));
                    } finally {
                        ThreadLocalLookupKey.clear();
//                        ShardInfoThreadLocal.clear();
                        latch.countDown();

//                        log.info(String.format("[ShardNoTxExecutor doInvoke] thread:%s method:%s shard:[%s,%s] delay:%d run_delay:%d", Thread.currentThread().getName(), method.getName(), shardInfo.getShardDB(), shardInfo.getShardTable(), System.currentTimeMillis() - now, System.currentTimeMillis() - _now));
                    }
                });
            } catch (Exception e) {
                res.put(shardInfo, ShardExecutor.Result.build(shardInfo, false, shardArgs, null, e, System.currentTimeMillis() - now));
                latch.countDown();
            }
        });

        //  等待所有的任务都完成，还没有完成的判定为超时
        try {
            latch.await(anno.timeout(), TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            throw new ShardDBRuntimeException(e);
        }
        shardInfoMap.forEach((shardInfo, shardArgs) -> {
            if (!res.containsKey(shardInfo)) {
                res.put(
                        shardInfo,
                        ShardExecutor.Result.build(
                                shardInfo,
                                false, shardArgs,
                                null,
                                new TimeoutException(String.format("Shard执行超时 run:(%dms) timeout:(%dms) method:%s class:%s db:%s table:%s", System.currentTimeMillis() - now, anno.timeout(), method.getName(), method.getDeclaringClass().getSimpleName(), shardInfo.getShardDB(), shardInfo.getShardTable())),
                                System.currentTimeMillis() - now));
            }
        });

        //  只要有1个失败，直接抛异常
        res.forEach(((shardInfo, r) -> {
            if (!r.isSuccess()) {
                throw new ShardDBRuntimeException(String.format("分库分表执行失败 shardDB:%s shardTable:%s method:%s class:%s", shardInfo.getShardDB(), shardInfo.getShardTable(), method.getName(), method.getDeclaringClass().getSimpleName()), r.getE());
            }
        }));

//        log.info(String.format("[ShardNoTxExecutor doInvoke] method:%s shardSize:%d delay:%d", method.getName(), shardInfoMap.size(), System.currentTimeMillis() - now));

        return res;
    }



    /**
     *  空值
     */
    protected Pagement NULL = new Pagement(){
        public Class<? extends Annotation> annotationType() {
            return null;
        }
        public String count() {
            return "";
        }
    };
    /**
     *  method与Pagement的映射
     */
    protected CachingKVTools<Method, Pagement> pagementMapper = new CachingKVTools<>(NULL, m -> {
        return m.getAnnotation(Pagement.class);
    });


    /**
     *  非事务执行方式
     */
    public static interface IShardNoTxDoInvoker {
        /**
         *  执行
         */
        ShardExecutor.Result doInvoke(Shard anno, ShardInfo shardInfo, Object[] args, Object proxy, Method method);
    }
}
